import cv2
import numpy as np
import os


def show_img(img, title="title"):
    cv2.imshow(title, img)
    # 关闭窗口
    cv2.waitKey()
    cv2.destroyAllWindows()


# img_a = 'temp/a1.jpg'
# img_a = 'temp/a4.jpg'
# img_a = 'temp/b2.jpg'
img_a = 'temp/b1.jpg'
# img_a = 'temp/b5.jpg'


img = cv2.imread(img_a, 1)
img = cv2.resize(img, (720, 480))
show_img(img, 'img')

blur = cv2.GaussianBlur(img, (5, 5), 5)
gray = cv2.cvtColor(blur,cv2.COLOR_BGR2GRAY)
edges = cv2.Canny(gray, 30, 60, apertureSize=3)

kernel = np.ones((7, 7), np.uint8)
img_dilate = cv2.dilate(edges, kernel)

# --- 直线检测
USE_HOUGH_P = 1
COLOR = (255, 0, 0)

img_black = np.zeros(img.shape, np.uint8)
# img_black = img.copy()      # 绘制背景
be_used_img = img_dilate.copy()  # 用来提取直线的图片
# be_used_img = edges.copy()     # 用来提取直线的图片


lines = cv2.HoughLinesP(be_used_img, 3, 1.0 * np.pi / 180, threshold=200, minLineLength=150,  maxLineGap=30)
assert lines is not None, 'lines为空!'

# LINE_LEN = 2000
for (x1, y1, x2, y2) in lines[:, 0]:
    print(x1, y1, ";", x2, y2)
    cv2.line(img_black, (x1, y1), (x2, y2), COLOR, 3)  # 画直线

print('--- lines.shape:', lines.shape)
show_img(img_black, 'USE_HOUGH_P')